from numpy import exp, zeros, float64, vectorize, array, where, empty, linalg
import datetime, time

class He3Analyzer:
    """ an object that contains information about a particular 3He cell, 
    including initial polarization and transmission and initial time.
    Member functions give polarization efficiency and polarization-correction
    matrix as a function of time """
    
    def __init__(self, P0 = '', T0 = '', Gamma = '', sigmanl='', t0='', t0_str = '', Psm='', Pf='', name ='', id_num = '0'):
        """ initialization parameters for cell:
            P0 = polarization at time t0
            T0 = transmission at t0
            Gamma = decay time constant for cell (hours)
            sigmanl = proportional to how much 3He in the beam path
            t0 = reference time (when initial measurements done)
            Psm = (supermirror + sample) polarization
            Pf = flipper polarization = (2*efficiency - 1) """
        
        self.P0 = float(P0)
        self.T0 = float(T0)
        self.Gamma = float(Gamma)
        self.sigmanl = float(sigmanl)
        self.t0_str = t0_str # datetime string, format =  '%a %b %d %H:%M:%S %Y'
        if t0_str == '': # we're working with a timestamp rather than a string time and date
            self.t0 = float(t0)
            self.t0_str = time.ctime(float(t0))
        else: # or we have a formatted date '%a %b %d %H:%M:%S %Y'
            self.t0 = time.mktime(time.strptime(self.t0_str,  '%a %b %d %H:%M:%S %Y'))
        self.start_datetime = datetime.datetime.strptime(self.t0_str, '%a %b %d %H:%M:%S %Y')
        self.Psm = float(Psm)
        self.Pf = float(Pf)
        self.name = name
        self.id_num = int(id_num)
        self.NTMatrix = zeros((4,4), dtype = float64)
        
    def getPolarization(self, t):
        """ t is the time since t0, in seconds """
        Gamma_seconds = self.Gamma * 60. * 60.
        P = self.P0 * exp(-t / Gamma_seconds )
        return P
        
    def getTransmission(self, t, aligned = True):
        """ returns the transmission of the analyzer for neutrons aligned 
        or anti-aligned with analyzer field (3He polarization varies with time) """
        P = self.getPolarization(t)
        if aligned:
            T = self.T0 * exp( -(1.0 - P) * self.sigmanl )
        else:
            T = self.T0 * exp( -(1.0 + P) * self.sigmanl )
        return T
        
    def getTransmissionVec(self, t, aligned = [True]):
        """ returns the transmission of the analyzer for neutrons aligned 
        or anti-aligned with analyzer field (3He polarization varies with time) """
        T = empty(aligned.shape)
        P = self.getPolarization(t)
        T[aligned] = self.T0 * exp( -(1.0 - P) * self.sigmanl )
        T[-aligned] = self.T0 * exp( -(1.0 + P) * self.sigmanl )
        return T
        
    def getNumberIncident(self, aligned = True, flipper_on = False):
        """ corrects for the polarization efficiency of the front-end components (not time-dependent) """
        if aligned:
            if flipper_on:
                N = (1.0 - self.Pf * self.Psm) / 2.0
            else:
                N = (1.0 + self.Psm) / 2.0
        else:
            if flipper_on:
                N = (1.0 + self.Pf * self.Psm) / 2.0
            else:
                N = (1.0 - self.Psm) / 2.0
        return N
        
    def getNumberIncidentVec(self, front_spin_up = [True], flipper_on=[False]):
        """ corrects for the polarization efficiency of the front-end components (not time-dependent) """
        N = empty(front_spin_up.shape)
        N[front_spin_up*flipper_on] = (1.0 - self.Pf * self.Psm) / 2.0
        N[front_spin_up*(-flipper_on)] = (1.0 + self.Psm) / 2.0
        N[(-front_spin_up)*flipper_on] = (1.0 + self.Pf * self.Psm) / 2.0
        N[(-front_spin_up)*(-flipper_on)] = (1.0 - self.Psm) / 2.0
        return N

    def getNTMatrix(self, t):
        """ creates matrix elements for the polarization-correction 
            this assumes the order of elements is Rup-up, Rup-down, Rdown-down, Rdown-up
            and for I: Iup-up, Iup-down, Idown-up, Idown-down   """
            
        flipper_on = array([
        [False, False, False, False],
        [False, False, False, False],
        [True, True, True, True],
        [True, True, True, True]
        ])
        # front_spin_up: True if neutrons are spin-up arriving at sample
        front_spin_up = array([
        [True, True, False, False],
        [True, True, False, False],
        [True, True, False, False],
        [True, True, False, False]
        ])
        # back_spin_up: True if neutrons are spin-up arriving at detector
        back_spin_up = array([
        [True, False, False, True],
        [True, False, False, True],
        [True, False, False, True],
        [True, False, False, True]
        ])
        # He3_up: True if He3 polarization is normal (not reversed)
        He3_up = array([
        [True, True, True, True],
        [False, False, False, False],
        [True, True, True, True],
        [False, False, False, False]
        ])
        # back_aligned: True if scattered neutron is parallel to 3He polarization
        #back_aligned = array([
        #[True, False, False, True],
        #[False, True, True, False],
        #[True, False, False, True],
        #[False, True, True, False]
        #])
        back_aligned = (back_spin_up * He3_up) + ((-back_spin_up) * (-He3_up))
        
        N = self.getNumberIncidentVec(front_spin_up, flipper_on)
        T = self.getTransmissionVec(t, back_aligned)

        return N*T
        
    def getNTRow(self, t, flipper_on = False, He3_up = True):
        flipper_on = array([flipper_on])
        He3_up = array([He3_up])
        # front_spin_up: True if neutrons are spin-up arriving at sample
        front_spin_up = array([
        [True, True, False, False]
        ])
        # back_spin_up: True if neutrons are spin-up arriving at detector
        back_spin_up = array([
        [True, False, False, True]
        ])
        back_aligned = (back_spin_up * He3_up) + ((-back_spin_up) * (-He3_up))
        
        N = self.getNumberIncidentVec(front_spin_up, flipper_on)
        T = self.getTransmissionVec(t, back_aligned)
        
        NT = N*T

        return N*T
        
    def getNTMasked(self, t, flipper_on_select = False, He3_up_select = True):
        flipper_on = array([
        [False, False, False, False],
        [False, False, False, False],
        [True, True, True, True],
        [True, True, True, True]
        ])
        # front_spin_up: True if neutrons are spin-up arriving at sample
        front_spin_up = array([
        [True, True, False, False],
        [True, True, False, False],
        [True, True, False, False],
        [True, True, False, False]
        ])
        # back_spin_up: True if neutrons are spin-up arriving at detector
        back_spin_up = array([
        [True, False, False, True],
        [True, False, False, True],
        [True, False, False, True],
        [True, False, False, True]
        ])
        # He3_up: True if He3 polarization is normal (not reversed)
        He3_up = array([
        [True, True, True, True],
        [False, False, False, False],
        [True, True, True, True],
        [False, False, False, False]
        ])
        
        back_aligned = (back_spin_up * He3_up) + ((-back_spin_up) * (-He3_up))
        
        N = self.getNumberIncidentVec(front_spin_up, flipper_on)
        T = self.getTransmissionVec(t, back_aligned)
        NT = zeros((4,4))
        mask = ((flipper_on == flipper_on_select) * (He3_up == He3_up_select))
        print mask
        NT[mask] = (N*T)[mask]

        return NT
    
